package com.tianji.service.impl;

import cn.hutool.core.date.DateUtil;
import com.tianji.common.utils.UserContext;
import com.tianji.config.SystemPromptConfig;
import com.tianji.entity.ChatMessage;
import com.tianji.entity.ChatSessionStatus;
import com.tianji.enums.ChatEventTypeEnum;
import com.tianji.enums.MessageStatusEnum;
import com.tianji.enums.MessageTypeEnum;
import com.tianji.service.ChatMessageService;
import com.tianji.service.ChatService;
import com.tianji.service.ChatSessionStatusService;
import com.tianji.vo.ChatEventVO;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Optional;

@Slf4j
@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final ChatClient chatClient;
    private final SystemPromptConfig systemPromptConfig;
    private final ChatMemory chatMemory;
    private final ChatMessageService chatMessageService;
    private final ChatSessionStatusService chatSessionStatusService;

    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        // 获得对话id
        var conversationId = ChatService.getConversationId(sessionId);
        Long userId = UserContext.getUser();

        // 保存用户问题到数据库
        saveUserMessage(sessionId, userId, conversationId, question);

        // 大模型输出内容的缓存器，用于在输出中断后的数据存储
        StringBuilder outputBuilder = new StringBuilder();

        return this.chatClient.prompt()
                .system(promptSystem -> promptSystem
                        .text(this.systemPromptConfig.getChatSystemMessage().get()) // 设置系统提示语
                        .param("now", DateUtil.now()) // 设置当前时间的参数
                )
                .advisors(advisor -> advisor.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId))
                .user(question)
                .stream()
                .chatResponse()
                .doFirst(() -> {  // 输出开始，标记正在输出
                    chatSessionStatusService.setGenerateStatus(sessionId, userId, true);
                })
                .doOnComplete(() -> { // 输出结束，清除标记
                    chatSessionStatusService.removeSessionStatus(sessionId, userId);
                })
                .doOnError(throwable -> chatSessionStatusService.removeSessionStatus(sessionId, userId)) // 错误时清除标记
                .doOnCancel(() -> {
                    // 当输出被取消时，保存输出的内容到历史记录中
                    this.saveStopHistoryRecord(conversationId, outputBuilder.toString());
                })
                // 输出过程中，判断是否正在输出，如果正在输出，则继续输出，否则结束输出
                .takeWhile(s -> Optional.ofNullable(chatSessionStatusService.getGenerateStatus(sessionId, userId)).orElse(false))
                .map(chatResponse -> {
                    // 获取大模型的输出的内容
                    String text = chatResponse.getResult().getOutput().getText();
                    // 追加到输出内容中
                    outputBuilder.append(text);
                    // 封装响应对象
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.just(ChatEventVO.builder()  // 标记输出结束
                        .eventType(ChatEventTypeEnum.STOP.getValue())
                        .build()));
    }

    @Override
    public void stop(String sessionId) {
        // 移除标记
        Long userId = UserContext.getUser();
        chatSessionStatusService.removeSessionStatus(sessionId, userId);
    }

    /**
     * 保存停止输出的记录
     *
     * @param conversationId 会话id
     * @param content        大模型输出的内容
     */
    private void saveStopHistoryRecord(String conversationId, String content) {
        // 保存到数据库
        String sessionId = conversationId.substring(conversationId.indexOf("_") + 1);
        Long userId = UserContext.getUser();
        saveAssistantMessage(sessionId, userId, conversationId, content, MessageStatusEnum.STOPPED);
        
        // 同时保存到内存
        this.chatMemory.add(conversationId, new AssistantMessage(content));
    }

    /**
     * 保存用户消息
     *
     * @param sessionId 会话ID
     * @param userId 用户ID
     * @param conversationId 对话ID
     * @param content 消息内容
     */
    private void saveUserMessage(String sessionId, Long userId, String conversationId, String content) {
        ChatMessage chatMessage = ChatMessage.builder()
                .sessionId(sessionId)
                .userId(userId)
                .conversationId(conversationId)
                .messageType(MessageTypeEnum.USER_MESSAGE.getValue())
                .content(content)
                .status(MessageStatusEnum.NORMAL.getValue())
                .build();
        chatMessageService.saveChatMessage(chatMessage);
    }

    /**
     * 保存助手消息
     *
     * @param sessionId 会话ID
     * @param userId 用户ID
     * @param conversationId 对话ID
     * @param content 消息内容
     * @param status 消息状态
     */
    private void saveAssistantMessage(String sessionId, Long userId, String conversationId, String content, MessageStatusEnum status) {
        ChatMessage chatMessage = ChatMessage.builder()
                .sessionId(sessionId)
                .userId(userId)
                .conversationId(conversationId)
                .messageType(MessageTypeEnum.ASSISTANT_MESSAGE.getValue())
                .content(content)
                .status(status.getValue())
                .build();
        chatMessageService.saveChatMessage(chatMessage);
    }
}
